package com.study.langchain4j.config.observability;

import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.UUID;

/**
 * @author xl-9527
 * @since 2025/10/28
 **/
@Component
public class ChatModelCustomObservability implements ChatModelListener {

    private static final Logger log = LoggerFactory.getLogger(ChatModelCustomObservability.class);

    @Override
    public void onRequest(final ChatModelRequestContext requestContext) {
        final Map<Object, Object> attributes = requestContext.attributes();
        final String traceId = UUID.randomUUID().toString().substring(0, 8);
        MDC.put("traceId", traceId);
        attributes.put("traceId", traceId);
        for (final Map.Entry<Object, Object> entry : attributes.entrySet()) {
            log.info("request {} -> {}", entry.getKey(), entry.getValue());
        }
    }

    @Override
    public void onResponse(final ChatModelResponseContext responseContext) {
        for (final Map.Entry<Object, Object> entry : responseContext.attributes().entrySet()) {
            log.info("response {} -> {}", entry.getKey(), entry.getValue());
        }
    }

    @Override
    public void onError(final ChatModelErrorContext errorContext) {
        log.error("model request error trace id -> {}", errorContext.attributes().get("traceId"), errorContext.error());
    }
}
